from langgraph.graph import END, StateGraph, START

from langgraph_coder.src.agent.model.graph_state import GraphState
from langgraph_coder.src.retriever.langchain_dataload import MyDataLoader

workflow = StateGraph(GraphState)

from langgraph_coder.src.agent.node.code_generate import GenerateNode
from langgraph_coder.src.agent.node.code_check import CheckNode
from langgraph_coder.src.agent.node.reflect import ReflectNode
from langgraph_coder.src.agent.edge.FinishDispatch import FinishDispatch

from langgraph_coder.src.agent.chain.generate_chain import GenerateChain
from langgraph.checkpoint.memory import InMemorySaver


class CodeGenerateAgent:
    def __init__(self, generate, check, reflect, finish, memory_inner):
        self.generate = generate
        self.check = check
        self.reflect = reflect
        self.finish = finish
        self.memory = memory_inner
        # 自动化加载
        self.graph = self._load()

    def _load(self):
        """
        构建工作流
        :return:
        """
        # 结点定义
        workflow.add_node("generate_node", self.generate)
        workflow.add_node("check_code", self.check)
        workflow.add_node("reflect_node", self.reflect)

        # 添加边
        workflow.add_edge(START, "generate_node")
        workflow.add_edge("generate_node", "check_code")
        # 条件边，必须经过跳转dispatch
        workflow.add_conditional_edges(
            "check_code",
            self.finish,
            {
                "end": END,
                "reflect": "reflect_node",
                "generate": "generate_node",
            },
        )
        workflow.add_edge("reflect_node", "generate_node")

        # 添加记忆，后续使用都要配置config
        return workflow.compile(checkpointer=memory)


if __name__ == '__main__':
    code_gen_chain = GenerateChain(MyDataLoader())

    # 结点定义
    generate_node = GenerateNode(code_gen_chain)
    check_node = CheckNode()
    reflect_node = ReflectNode(code_gen_chain)

    # 边定义
    finish_dispatch = FinishDispatch()

    # 内存记忆定义
    memory = InMemorySaver()

    # 构建实体
    agent = CodeGenerateAgent(generate_node, check_node, reflect_node, finish_dispatch, memory)

    # 配置memory的config
    config = {"configurable": {"thread_id": "1"}}

    question = "如何直接将字符串传递给可运行对象，并使用它来构建我提示所需的输入？?"
    solution = agent.graph.invoke({"messages": [("user", question)], "iterations": 0, "error": "", "max_iterations": 5, "flag": "do not reflect"}, config=config)
    print(solution["generation"])